#!/usr/bin/env python
import json
import sys
import carla
import numpy as np
import time
import random
from queue import Queue, Empty
import os
import copy
import logging

import pandas as pd
from carla import Transform, Rotation, Location
from pyproj import Proj
import math
import imageio
import threading

try:
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + '/carla')
except IndexError:
    pass

def points2pcd(PCD_FILE_PATH, points):
    '''
    :param PCD_FILE_PATH: pcd文件详细保存地址
    :param points: 待保存的点云数据
    :return:
    '''
    # 存放路径
    if os.path.exists(PCD_FILE_PATH):
        os.remove(PCD_FILE_PATH)
    # 写文件句柄
    handle = open(PCD_FILE_PATH, 'a')
    # 得到点云点数
    point_num = points.shape[0]
    # pcd头部（重要）
    handle.write(
        '# .PCD v0.7 - Point Cloud Data file format\nVERSION 0.7\nFIELDS x y z intensity\nSIZE 4 4 4 4\nTYPE F F F F\nCOUNT 1 1 1 1')
    string = '\nWIDTH ' + str(point_num)
    handle.write(string)
    handle.write('\nHEIGHT 1\nVIEWPOINT 0 0 0 1 0 0 0')
    string = '\nPOINTS ' + str(point_num)
    handle.write(string)
    handle.write('\nDATA ascii')

    # 依次写入点
    for i in range(point_num):
        string = '\n' + str(points[i, 0]) + ' ' + str(points[i, 1]) + ' ' + str(points[i, 2]) + ' ' + str(points[i, 3])
        handle.write(string)
    handle.close()

def sensor_callback(sensor_data, sensor_queue, sensor_name):
    '''
    处理sensor_data数据，比如将其保存到磁盘，然后只需要添加到队列中
    :param sensor_data: sensor_data数据
    :param sensor_queue: sensor_data数据队列
    :param sensor_name: sensor_data数据名称
    :return:
    '''
    sensor_queue.put((sensor_data.frame, sensor_data.timestamp, sensor_name, sensor_data))

def _parse_lidar_cb(lidar_data):
    '''
    处理lidar数据
    :param lidar_data: 待处理的lidar数据
    :return: lidar数据队列
    '''
    points = np.frombuffer(lidar_data.raw_data, dtype=np.dtype('f4'))
    points = copy.deepcopy(points)
    points = np.reshape(points, (int(points.shape[0] / 4), 4))
    return points

def _parse_image_cb(image):
    '''
    处理camera数据
    :param image: 待处理的camera数据
    :return: camera数据队列
    '''
    array = np.frombuffer(image.raw_data, dtype=np.dtype("uint8"))
    array = np.reshape(array, (image.height, image.width, 4))
    array = array[:, :, :3]
    array = array[:, :, ::-1]
    return array

def add_camera(world, number, camera_trans, width, height, fov, tick):
    '''
    添加相机传感器
    :param world: 世界
    :param number: 第几个相机传感器
    :param camera_trans: 相机的放置位置
    :param width:  图像像素
    :param height: 图像像素
    :param fov: 相机水平视场角
    :param tick:  相机的图像捕捉周期
    '''
    blueprint_library = world.get_blueprint_library()
    camera_bp = blueprint_library.find('sensor.camera.rgb')
    camera_bp.set_attribute("image_size_x", str(width))
    camera_bp.set_attribute("image_size_y", str(height))
    camera_bp.set_attribute("fov", str(fov))
    camera_bp.set_attribute('sensor_tick', str(tick))
    camera = world.spawn_actor(camera_bp, camera_trans, attach_to=None)
    camera.listen(lambda data: sensor_callback(data, sensor_queue, "camera_{}".format(number)))
    sensor_list.append(camera)

def add_732L_lidar(world, lidar_trans, number):
    '''
    添加万集732L雷达传感器
    :param world: 世界
    :param lidar_trans: 雷达的放置位置
    '''

    blueprint_library = world.get_blueprint_library()

    lidar_bp_16_1 = blueprint_library.find('sensor.lidar.ray_cast')
    lidar_bp_8_1 = blueprint_library.find('sensor.lidar.ray_cast')
    lidar_bp_4_1 = blueprint_library.find('sensor.lidar.ray_cast')
    lidar_bp_4_2 = blueprint_library.find('sensor.lidar.ray_cast')

    lidar_bp_16_1.set_attribute('channels', '16')
    lidar_bp_16_1.set_attribute('upper_fov', '0')
    lidar_bp_16_1.set_attribute('lower_fov', '-9')
    lidar_bp_16_1.set_attribute('points_per_second', '288000')
    lidar_bp_16_1.set_attribute('range', '200')
    lidar_bp_16_1.set_attribute('rotation_frequency', '10')
    lidar_bp_16_1.set_attribute('atmosphere_attenuation_rate', '0.004')
    lidar_bp_16_1.set_attribute('dropoff_general_rate', '0.0')
    lidar_bp_16_1.set_attribute('dropoff_intensity_limit', '0.0')
    lidar_bp_16_1.set_attribute('dropoff_zero_intensity', '0.0')
    lidar_bp_16_1.set_attribute('noise_seed', '0.0')
    lidar_bp_16_1.set_attribute('noise_stddev', '0.0')

    lidar_bp_8_1.set_attribute('channels', '8')
    lidar_bp_8_1.set_attribute('upper_fov', '-10')
    lidar_bp_8_1.set_attribute('lower_fov', '-17')
    lidar_bp_8_1.set_attribute('points_per_second', '144000')
    lidar_bp_8_1.set_attribute('range', '200')
    lidar_bp_8_1.set_attribute('rotation_frequency', '10')
    lidar_bp_8_1.set_attribute('atmosphere_attenuation_rate', '0.004')
    lidar_bp_8_1.set_attribute('dropoff_general_rate', '0.0')
    lidar_bp_8_1.set_attribute('dropoff_intensity_limit', '0.0')
    lidar_bp_8_1.set_attribute('dropoff_zero_intensity', '0.0')
    lidar_bp_8_1.set_attribute('noise_seed', '0.0')
    lidar_bp_8_1.set_attribute('noise_stddev', '0.0')

    lidar_bp_4_1.set_attribute('channels', '4')
    lidar_bp_4_1.set_attribute('upper_fov', '-19')
    lidar_bp_4_1.set_attribute('lower_fov', '-25')
    lidar_bp_4_1.set_attribute('points_per_second', '72000')
    lidar_bp_4_1.set_attribute('range', '200')
    lidar_bp_4_1.set_attribute('rotation_frequency', '10')
    lidar_bp_4_1.set_attribute('atmosphere_attenuation_rate', '0.004')
    lidar_bp_4_1.set_attribute('dropoff_general_rate', '0.0')
    lidar_bp_4_1.set_attribute('dropoff_intensity_limit', '0.0')
    lidar_bp_4_1.set_attribute('dropoff_zero_intensity', '0.0')
    lidar_bp_4_1.set_attribute('noise_seed', '0.0')
    lidar_bp_4_1.set_attribute('noise_stddev', '0.0')

    lidar_bp_4_2.set_attribute('channels', '4')
    lidar_bp_4_2.set_attribute('upper_fov', '-28')
    lidar_bp_4_2.set_attribute('lower_fov', '-37')
    lidar_bp_4_2.set_attribute('points_per_second', '72000')
    lidar_bp_4_2.set_attribute('range', '200')
    lidar_bp_4_2.set_attribute('rotation_frequency', '10')
    lidar_bp_4_2.set_attribute('atmosphere_attenuation_rate', '0.004')
    lidar_bp_4_2.set_attribute('dropoff_general_rate', '0.0')
    lidar_bp_4_2.set_attribute('dropoff_intensity_limit', '0.0')
    lidar_bp_4_2.set_attribute('dropoff_zero_intensity', '0.0')
    lidar_bp_4_2.set_attribute('noise_seed', '0.0')
    lidar_bp_4_2.set_attribute('noise_stddev', '0.0')

    lidar_16_1 = world.spawn_actor(lidar_bp_16_1, lidar_trans, attach_to=None)
    lidar_8_1 = world.spawn_actor(lidar_bp_8_1, lidar_trans, attach_to=None)
    lidar_4_1 = world.spawn_actor(lidar_bp_4_1, lidar_trans, attach_to=None)
    lidar_4_2 = world.spawn_actor(lidar_bp_4_2, lidar_trans, attach_to=None)
    lidar_16_1.listen(lambda data: sensor_callback(data, sensor_queue, "lidar_16_1_{}".format(number)))
    lidar_8_1.listen(lambda data: sensor_callback(data, sensor_queue, "lidar_8_1_{}".format(number)))
    lidar_4_1.listen(lambda data: sensor_callback(data, sensor_queue, "lidar_4_1_{}".format(number)))
    lidar_4_2.listen(lambda data: sensor_callback(data, sensor_queue, "lidar_4_2_{}".format(number)))
    sensor_list.append(lidar_16_1)
    sensor_list.append(lidar_8_1)
    sensor_list.append(lidar_4_1)
    sensor_list.append(lidar_4_2)

def save_data(w_timestamp, concat_points_list, camera_data_list):
    '''
    lidar/camera等数据的实时保存
    :param w_timestamp: 以实时时间信息为保存的数据的命名
    :param concat_points_list: 待保存的lidar数据
    :param camera_data_list: 待保存的camera数据
    :return:
    '''
    try:
        # 实时保存点云数据(pcd)
        if SAVE_PCD:
            for x in range(Lidar_Number):
                pcd_path_floder = save_lidar_path + "/" + str(x)
                if not os.path.exists(pcd_path_floder):
                    os.makedirs(pcd_path_floder)
                pcd_path = os.path.join(pcd_path_floder, "%s.pcd" % str(w_timestamp))
                # points = pd.DataFrame(concat_points_list[x], columns=['x', 'y', 'z', 'i'])
                # points.to_csv(pcd_path, index=False)
                points2pcd(pcd_path, concat_points_list[x])

        # # 实时保存图像
        # if SAVE_Camera:
        #     for i in range(Camera_Number):
        #         img_path_floder = save_img_path + "/" + str(i)
        #         if not os.path.exists(img_path_floder):
        #             os.makedirs(img_path_floder)
        #         img_path = os.path.join(img_path_floder, "%s.jpg" % str(w_timestamp))
        #         # print(type(img_path), ", ", type(camera_data_list[i]))
        #         camera_data_list[i].save_to_disk(img_path)
    except Exception as e:
        print(e)

def trajectory(world, vehicles_list, w_frame):
    result = None
    car_list = list()

    all_vehicle_actors = world.get_actors(vehicles_list)
    if len(all_vehicle_actors) > 0:
        for vehicle in all_vehicle_actors:
            vehicle_id = vehicle.id
            vehicle_type = vehicle.type_id
            x = vehicle.get_transform().location.x
            y = vehicle.get_transform().location.y
            vehicle_angle = vehicle.get_transform().rotation.yaw
            v = vehicle.get_velocity()
            vehicle_speed = 3.6 * math.sqrt(v.x ** 2 + v.y ** 2 + v.z ** 2)
            length = vehicle.bounding_box.extent.x * 2.0
            width = vehicle.bounding_box.extent.y * 2.0
            height = vehicle.bounding_box.extent.z * 2.0
            longitude, latitude = rule(x, -y, inverse=True)
            # dis = math.sqrt((x-x0) ** 2 + (y-y0) ** 2)
            # if x == 0 and y == 0:
            #     vehicle.destroy()
            #     vehicles_list.remove(vehicle_id)
            # else:
            car = {"id": vehicle_id, "type": vehicle_type, "longitude": longitude, "latitude": latitude,
                   "angle": vehicle_angle, "speed": vehicle_speed, "length": length, "width": width, "height": height}
            car_list.append(car)
    # if len(car_list) > 0:
    result = {"participants": car_list, "participantNum": len(car_list),
              "frameID": str(w_frame), "globalTimeStamp": str(time.time())}
    return result

def main():
    #  ======================================跟服务器实现连接=============================================
    client = carla.Client("localhost", 2000)
    client.set_timeout(20.0)  # 设置这个客户端连接超时时间
    # print(client.get_available_maps())  # 获取可用的地图列表（自带的地图，无自定义地图）
    # world = client.load_world('SDWR0525')
    world = client.get_world()  # 如果在UE4中运行场景，则使用该world
    synchronous_master = False

    try:
        # =================================调整路网中的环境可见度======================================
        # 使得道路中自动生成的交通灯为不可见状态
        env_objs1 = world.get_environment_objects(carla.CityObjectLabel.TrafficLight)
        env_objs2 = world.get_environment_objects(carla.CityObjectLabel.Poles)

        env_list = []
        for env_obj in env_objs1:
            env_list.append(env_obj.id)
        for env_obj in env_objs2:
            env_list.append(env_obj.id)
        world.enable_environment_objects(env_list, False)

        # ====================================获取世界视角=============================================
        spectator = world.get_spectator()
        spectator.set_transform(view_transform)
        world.debug.draw_string(camera_trans2.location, 'X', draw_shadow=False, color=carla.Color(r=255, g=0, b=0),
                                life_time=500, persistent_lines=True)

        # ============================== 修改世界设置:更新步长和同步模式===================================
        # 让车辆按照交通规则在世界中行驶
        # 设置TM接入server的端口，默认8000
        traffic_manager = client.get_trafficmanager(8000)
        # TM里的每一辆车至少保持2米的安全距离
        # traffic_manager.set_global_distance_to_leading_vehicle(2.0)
        # TM里每一辆车都是默认速度的40%（默认所有车辆限速30km/h）
        # traffic_manager.global_percentage_speed_difference(40.0)
        if True:
            settings = world.get_settings()
            traffic_manager.set_synchronous_mode(True)
            if not settings.synchronous_mode:
                print("开启同步模式")
                synchronous_master = True
                settings.synchronous_mode = True
                settings.fixed_delta_seconds = 0.1
                world.apply_settings(settings)
            else:
                synchronous_master = False
                settings.fixed_delta_seconds = 0.1
                world.apply_settings(settings)

        # 控制世界的天气和时间（太阳的位置） 万里无云，没有降雨，太阳的角度为90
        weather = carla.WeatherParameters(sun_altitude_angle=90.0)
        world.set_weather(weather)

        # =============================传感器=================================
        for i in range(Camera_Number):
            camera_trans = camera_trans_list[i]
            add_camera(world, i, camera_trans, IM_WIDTH, IM_HEIGHT, fov=90, tick=0.1)
        for i in range(Lidar_Number):
            lidar_trans = lidar_trans_list[i]
            add_732L_lidar(world, lidar_trans, i)

        with open(realdata) as sim_file:
            sim_data_dic = json.load(sim_file)
        real_iterator = sim_data_dic.__iter__()

        while True:

            world.tick()
            w_frame = world.get_snapshot().frame
            p_timestamp = world.get_snapshot().timestamp.platform_timestamp
            print("\nWorld's frame:{0}, time: {1}".format(w_frame, p_timestamp))

            try:
                errorids = []
                try:
                    timestep = real_iterator.__next__()
                except StopIteration:
                    break
                nowids = list(sim_data_dic[timestep].keys())
                # print(len(nowids))

                # 添加交通流
                for k, v in sim_data_dic[timestep].items():
                    if v is not None:
                        loc = v['Location']
                        rot = v['Rotation']
                        trans = Transform(carla.Location(loc[0], loc[1], loc[2]),
                                          carla.Rotation(rot[0], rot[1], rot[2]))
                        vehtype = v['type']
                        if k not in spawed_ids.keys():
                            bp1 = world.get_blueprint_library().find(vehtype)
                            batch = [
                                carla.command.SpawnActor(bp1, trans).then(
                                    carla.command.SetSimulatePhysics(carla.command.FutureActor, True))
                            ]
                            response = client.apply_batch_sync(batch, False)[0]

                            if response.error:
                                errorids.append(k)
                            else:
                                spawed_ids[k] = response.actor_id
                                vehicles_id_list.append(response.actor_id)
                        else:
                            # 更新当前帧位姿信息
                            carlaid = spawed_ids[k]
                            vehicle = world.get_actor(carlaid)
                            if Correction_yaw:
                                # 获取车辆所在位置附近的道路点
                                v_waypoint = world.get_map().get_waypoint(
                                    vehicle.get_location(), project_to_road=True,
                                    lane_type=(carla.LaneType.Driving))
                                yaw = v_waypoint.transform.rotation.yaw
                                inside_junction = v_waypoint.is_junction
                                # inside_junction = False
                                if inside_junction:
                                    dx = trans.location.x - vehicle.get_location().x
                                    dy = trans.location.y - vehicle.get_location().y
                                    if dx != 0.0:
                                        yaw = math.atan2(dy, dx) * 180.0 / math.pi
                                    else:
                                        yaw = vehicle.get_transform().rotation.yaw
                                trans = Transform(carla.Location(loc[0], loc[1], loc[2]),
                                                  carla.Rotation(rot[0], yaw, rot[2]))

                            vehicle.set_transform(trans)
                        # if vehicle is not None:
                        #     vehicle.set_transform(trans)
                print("已有：", len(spawed_ids))

                destodyed_ids = [id for id in spawed_ids if id not in nowids]
                for did in destodyed_ids:
                    carlaid = spawed_ids[did]
                    vehicle = world.get_actor(carlaid)

                    if vehicle is not None:
                        vehicle.destroy()
                    del spawed_ids[did]
                print("更新：", len(nowids), "销毁:", len(destodyed_ids))
                print("错误：", len(errorids))

                # 获取车辆轨迹
                vehicle_list = vehicles_id_list
                result = trajectory(world, vehicle_list, w_frame)
                # 保存车辆轨迹
                if not os.path.exists(save_trajectory_path):
                    os.makedirs(save_trajectory_path)
                json_name = os.path.join(save_trajectory_path, "%s.json" % str(0))
                with open(json_name, 'a') as file:
                    json.dump(result, file)
                    file.write('\n')

                lidar_data_list = []
                concat_points_list = []
                camera_data_list = []
                for n in range(Lidar_Number):
                    lidar_data_list.append([])
                    concat_points_list.append([])
                for p in range(Camera_Number):
                    camera_data_list.append([])
                # 数据存储
                for i in range(0, len(sensor_list)):
                    s_frame, s_timestamp, s_name, s_data = sensor_queue.get(True, 1.0)
                    s_name_list = s_name.split('_')
                    sensor_type = str(s_name_list[0])
                    sensor_id = int(s_name_list[len(s_name_list)-1])
                    if sensor_type == "lidar":
                        lidar_data_list[sensor_id].append(_parse_lidar_cb(s_data))
                    elif s_name == "camera_0":
                        img0_path_floder = save_img_path + "/" + str(0)
                        if not os.path.exists(img0_path_floder):
                            os.makedirs(img0_path_floder)
                        img0_path = os.path.join(img0_path_floder, "%s.jpg" % str(w_frame))
                        s_data.save_to_disk(img0_path)
                    elif s_name == "camera_1":
                        img1_path_floder = save_img_path + "/" + str(1)
                        if not os.path.exists(img1_path_floder):
                            os.makedirs(img1_path_floder)
                        img1_path = os.path.join(img1_path_floder, "%s.jpg" % str(w_frame))
                        s_data.save_to_disk(img1_path)
                # lidar数据处理
                for m in range(Lidar_Number):
                    concat_points_list[m] = np.concatenate(lidar_data_list[m], axis=0)
                    concat_points_list[m][:, 1] = [-p for p in concat_points_list[m][:, 1]]

                # save_queue.put((concat_points_list[0], w_frame))
                # 保存传感器数据
                save_data(w_frame, concat_points_list, camera_data_list)

            except Empty:
                print("[Warning] Some of the sensor information is missed")

    finally:
        if synchronous_master:
            settings = world.get_settings()
            settings.synchronous_mode = False
            settings.fixed_delta_seconds = None
            world.apply_settings(settings)
        for sensor in sensor_list:
            sensor.destroy()
        client.apply_batch([carla.command.DestroyActor(x) for x in actor_list])
        print("All cleaned up!")

if __name__ == "__main__":

    rule = Proj("+proj=tmerc +lat_0=36.63498592769697 +lon_0=117.5620618908321 +ellps=WGS84 ")

    local_time = time.localtime(time.time())
    timestamp = time.strftime("%Y%m%d_%H%M%S", local_time)
    save_lidar_path = '/home/gj/Carla_Output/lidar/' + str(timestamp)
    save_img_path = '/home/gj/Carla_Output/camera/' + str(timestamp)
    save_trajectory_path = '/home/gj/Carla_Output/trajectory/' + str(timestamp)
    realdata = "/home/gj/Carla_0.9.13/PythonAPI/work/tessng_track2.json"

    sensor_queue = Queue()
    actor_list, sensor_list = [], []
    camera_trans_list, lidar_trans_list = [], []
    spawed_ids = {}
    vehicles_id_list = []
    save_threads = []
    save_queue = Queue()

    Camera_Number = 2
    Lidar_Number = 1

    Correction_yaw = False  # 是否修正跟踪车辆的航向角，根据车辆所在位置最近的waypoint的航向角信息修正
    SAVE_PCD = True
    SAVE_Camera = True

    IM_WIDTH = 1920
    IM_HEIGHT = 1080

    # 视角、雷达、相机的位置和姿态
    # x0, y0 = rule(117.595160, 36.639694)
    x0, y0 = rule(117.56610262, 36.62692877352896)

    view_transform = Transform(Location(x=x0, y=-y0, z=6.0), Rotation(pitch=-15, yaw=120, roll=0))

    camera_trans1 = Transform(Location(x=x0, y=-y0, z=6.0), Rotation(pitch=-15.0, yaw=160.0, roll=0.0))
    camera_trans2 = Transform(Location(x=x0, y=-y0, z=6.0), Rotation(pitch=-15.0, yaw=20.0, roll=0.0))
    lidar_trans1 = Transform(Location(x=x0, y=-y0, z=6.0), Rotation(pitch=0.0, yaw=173.5, roll=0.0))

    camera_trans_list.append(camera_trans1)
    camera_trans_list.append(camera_trans2)
    lidar_trans_list.append(lidar_trans1)

    try:
        main()
    except KeyboardInterrupt:
        print(' - Exited by user.')
